import pandas as pd

from app_config import get_engine_ts
from app_config import get_pro
import os

from datetime import datetime
from dateutil.relativedelta import relativedelta


""" 科创50
"""
def offset_date_6M(original_date_str):
    # 将输入的日期字符串转换为日期对象
    original_date = datetime.strptime(original_date_str, '%Y%m%d')

    # 计算6个月后的日期
    new_date = original_date + relativedelta(months=6)

    # 将结果日期格式化为 'yyyyMMdd' 字符串
    return new_date.strftime('%Y%m%d')


def calc_kc50(strStartDate='', strEndDate='', str_filePre='', str_000688Date=''):
    if not os.path.exists(str_filePre):
        os.makedirs(str_filePre)
        print(f"文件夹 '{str_filePre}' 创建成功")
    engine = get_engine_ts()
    """     1 样本空间    """
    ''' i. 上市时间超过 6 个月；待科创板上市满 12 个月的证券数量达 100 只至
    150 只后，上市时间调整为超过 12 个月；
    '''
    kcb = get_pro().stock_basic(market='科创板', list_status='L', fields='ts_code,symbol,name,area,industry,list_date')

    # 筛选 list_date 小于 '20230601' 的数据 一年前
    #  下: 上市时间超过 6 个月
    offset_6m = offset_date_6M(strStartDate)

    print("时间 : =================")
    print(f"start date: {strStartDate}")
    print(f"6M date: {offset_6m}")
    print(f"end date: {strEndDate}")

    kc_6_mouth = kcb[kcb['list_date'] <= offset_6m]
    # kcOneYear = kcb[kcb['list_date'] < strEndDate]
    """     2、选样方法    """
    '''2.1对样本空间内的证券按照过去一年的日均成交金额由高到低排名，剔除排名后 10%的证券作为待选样本；'''
    # 转换 ts_code 列为元组，并生成 SQL 中 IN 子句所需的格式
    ts_codes = kc_6_mouth['ts_code'].tolist()

    # 构建查询语句
    query = f"""
    SELECT ts_code,trade_date, amount FROM `daily`
    WHERE ts_code IN ({','.join(f"'{code}'" for code in ts_codes)})
    AND trade_date >= '{strStartDate}'
    AND trade_date <= '{strEndDate}'
    """

    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['amount'] = result_df['amount'] / 1000
    grouped_df = result_df.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', '日均成交额']
    merge = pd.merge(kc_6_mouth, grouped_df, on='ts_code', how='left')

    # 构建查询语句
    query = f"""
        SELECT ts_code,trade_date, total_mv FROM `daily_basic`
        WHERE ts_code IN ({','.join(f"'{code}'" for code in ts_codes)})
        AND trade_date >= '{strStartDate}'
        AND trade_date <= '{strEndDate}'
        """
    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['total_mv'] = result_df['total_mv'] / 10000
    grouped_df_total_mv = result_df.groupby('ts_code')['total_mv'].mean().reset_index()
    grouped_df_total_mv.columns = ['ts_code', '日均总市值']
    pd_merge = pd.merge(merge, grouped_df_total_mv, on='ts_code', how='left')

    # 按 日均成交额 列降序排序
    df_sorted = pd_merge.sort_values(by='日均成交额', ascending=False).reset_index(drop=True)
    df_sorted['i_avgAmount'] = df_sorted.index
    # 获取前 90% 的数据
    num_rows = len(df_sorted)
    top_50_percent = df_sorted.head(int(num_rows * 0.9))

    top_50_percent = top_50_percent.sort_values('日均总市值', ascending=False).reset_index(drop=True)
    top_50_percent['i_totalMv'] = top_50_percent.index
    df_result_kc50 = top_50_percent.head(120)
    # df_result_kc50.to_excel(str_filePre + '科创50.xlsx', index=True)

    # =======================
    # ========表格合并========
    # =======================

    # if str_000688Date == '':
    #     print("程序结束")
    #     sys.exit()
    #
    # print("开始合并")
    df_930955 = get_pro().index_weight(index_code='000688.SH', start_date=str_000688Date, end_date=str_000688Date)
    df_930955.rename(columns={'con_code': 'ts_code'}, inplace=True)
    df_930955.rename(columns={'trade_date': 'list_date'}, inplace=True)
    df_930955['symbol'] = df_930955['ts_code'].str.split('.').str[0]
    df_930955.drop(columns=['weight'], inplace=True)
    df_930955.to_excel(str_filePre + str_000688Date + '-000688.xlsx')

    # 获取 dataframe1 中的所有列
    columns_dataframe1 = df_result_kc50.columns

    # 创建缺失列并填充为 '*'
    for column in columns_dataframe1:
        if column not in df_930955.columns:
            df_930955[column] = '**'

    # 将 dataframe2 的列顺序调整为与 dataframe1 一致
    df_930955 = df_930955[columns_dataframe1]

    # 追加 dataframe2 到 dataframe1
    result = pd.concat([df_result_kc50, df_930955], ignore_index=True)

    result.to_excel(str_filePre + '科创50_6M.xlsx')


if __name__ == '__main__':
    """2024年四季度"""
    strStartDate = '20231101'
    strEndDate = '20241031'
    str_filePre = 'zfile/科创50_2024季度4/'
    str_000688Date = '20240930'
    calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)

    # """2024年三季度"""
    # strStartDate = '20230801'
    # strEndDate = '20240731'
    # str_filePre = 'zfile/科创50_2024季度3/'
    # str_000688Date = '20240628'
    # calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)

    # # """2024年2季度"""
    # strStartDate = '20230501'
    # strEndDate = '20240430'
    # str_filePre = 'zfile/科创50_2024季度2/'
    # str_000688Date = '20240628'
    # calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)
    #
    # # """2024年1季度"""
    # strStartDate = '20230201'
    # strEndDate = '20240131'
    # str_filePre = 'zfile/科创50_2024季度1/'
    # str_000688Date = '20240430'
    # calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)
    #
    # # """2023年四季度"""
    # strStartDate = '20221101'
    # strEndDate = '20231031'
    # str_filePre = 'zfile/科创50_2023季度4/'
    # str_000688Date = '20240131'
    # calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)
    #
    # # """2023年三季度"""
    # strStartDate = '2022801'
    # strEndDate = '20230731'
    # str_filePre = 'zfile/科创50_2023季度3/'
    # str_000688Date = '20231031'
    # calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)
    #
    # # """2023年二季度"""
    # strStartDate = '2022501'
    # strEndDate = '20230430'
    # str_filePre = 'zfile/科创50_2023季度2/'
    # str_000688Date = '20230731'
    # calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)
    #
    # # """2023年一季度"""
    # strStartDate = '20220201'
    # strEndDate = '20230131'
    # str_filePre = 'zfile/科创50_2023季度1/'
    # str_000688Date = '20230428'
    # calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)
    #
    # # """2022年四季度"""
    # strStartDate = '20211101'
    # strEndDate = '20221031'
    # str_filePre = 'zfile/科创50_2022季度4/'
    # str_000688Date = '20230131'
    # calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)
    #
    # # """2022年三季度"""
    # strStartDate = '20210801'
    # strEndDate = '20220731'
    # str_filePre = 'zfile/科创50_2022季度3/'
    # str_000688Date = '20221031'
    # calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)
    #
